import torch
import torch.nn as nn
import torch.nn.functional as F

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class GRN(nn.Module):

    def __init__(self, dim):
        super().__init__()
        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))

    def forward(self, x):
        Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
        return self.gamma * (x * Nx) + self.beta + x


class ConvNeXtV2Block(nn.Module):

    def __init__(self, dim, layer_scale_init_value=1e-6):
        super().__init__()
        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
        self.norm = nn.LayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(dim, 4 * dim)
        self.grn = GRN(4 * dim)
        self.pwconv2 = nn.Linear(4 * dim, dim)
        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(dim)) if layer_scale_init_value > 0 else None

    def forward(self, x):
        input = x
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.grn(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.permute(0, 3, 1, 2)
        return input + x

class SpatialAttention(nn.Module):
    def __init__(self,in_channels,kernel_size=7):
        super(SpatialAttention, self).__init__()
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()
        self.conv2 =nn.Sequential(nn.Conv2d(in_channels, in_channels, kernel_size, padding=padding, bias=False),
                                    nn.BatchNorm2d(in_channels),
                                   nn.ReLU()
                                  )

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x_cat = torch.cat([avg_out, max_out], dim=1)
        spatial_weights = self.sigmoid(self.conv1(x_cat))
        return self.conv2(x * spatial_weights)

class MultiStreamConvNeXtV2(nn.Module):

    def __init__(self, in_channels=3, num_classes=5):
        super().__init__()

        # 分支1（原始图像处理）
        self.stem1 = nn.Sequential(
            nn.Conv2d(in_channels, 32, 3, padding=1, bias=False),
            nn.MaxPool2d(2, 2)
        )
        self.block1 = nn.Sequential(
            ConvNeXtV2Block(32),
            ConvNeXtV2Block(32),
            nn.MaxPool2d(2, 2)
        )

        self.stem2 = nn.Sequential(
            nn.Conv2d(in_channels, 32, 7, padding=3, bias=False),
            nn.MaxPool2d(2, 2)
        )
        self.block2 = nn.Sequential(
            ConvNeXtV2Block(32),
            ConvNeXtV2Block(32),
            nn.MaxPool2d(2, 2)
        )


        self.fc_block = nn.Sequential(
            nn.Linear(42, 128),
            nn.BatchNorm1d(128),  # 替换 Dropout 为 BatchNorm1d
            nn.GELU(),
            nn.Linear(128, 256),
            nn.BatchNorm1d(256),  # 替换 Dropout 为 BatchNorm1d
            nn.GELU(),
            nn.Linear(256, 512),
            nn.BatchNorm1d(512)  # 替换 Dropout 为 BatchNorm1d
        )

        self.class_head = nn.Sequential(
            nn.Linear(1024, 1024),
            nn.BatchNorm1d(1024),
            nn.GELU(),
            nn.Linear(1024, 512),
            nn.BatchNorm1d(512),
            nn.GELU(),
            nn.Linear(512, num_classes)
        )

        # 特征融合模块
        self.fusion = nn.ModuleList([
            ConvNeXtV2Block(64),
            ConvNeXtV2Block(128),
            ConvNeXtV2Block(256)
        ])

        # 分类头（ConvNeXtV2设计规范）
        self.before_add = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.LayerNorm(512),
        )

        self.channel_adjust = nn.ModuleDict({
            'layer1': nn.Conv2d(64, 128, 1),
            'layer2': nn.Conv2d(128, 256, 1),
            'layer3': nn.Conv2d(256, 512, 1)
        })

    def forward(self, x1, x2, x3):

        f1 = self.stem1(x1)
        f1 = self.block1(f1)


        f2 = self.stem2(x2)
        f2 = self.block2(f2)


        fused = torch.cat([f1, f2], dim=1)  # 通道拼接

        # space1 = SpatialAttention(64).to(device=device)
        # fused =space1(fused)
        # 多级特征融合
        for i, block in enumerate(self.fusion):
            fused = block(fused)
            fused = self.channel_adjust[f'layer{i + 1}'](fused)
            fused = F.max_pool2d(fused, 2)

        # space2 = SpatialAttention(512).to(device=device)
        # fused = space2(fused)

        tabular_feat = self.fc_block(x3)

        # 最终分类
        before_img = self.before_add(fused)
        out = torch.cat([before_img, tabular_feat], dim=1)
        out = self.class_head(out)  # 特征融合

        return out

from torchinfo import summary

if __name__ == "__main__":
    model = MultiStreamConvNeXtV2(in_channels=3, num_classes=4).to(device=device)
    # rgb = torch.randn(2, 3, 256, 256)
    # depth = torch.randn(2, 3, 256, 256)
    # tabular = torch.randn(2, 42)
    # output = model(rgb, depth, tabular)


    summary(
        model,
        input_size=[(2,3, 256, 256), (2,3, 256, 256), (2,42)],  # 对应三个输入分支
        depth=5,  # 显示嵌套层级
        col_names=["input_size", "output_size", "num_params"]
    )